import os
from typing import Tuple

try:
    import folder_paths
except Exception:
    folder_paths = None


class YouTubeDownloader:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "url": ("STRING", {"default": "", "multiline": False}),
                "mode": (["video", "audio", "subtitles"], {"default": "video"}),
                "quality": (["highest", "144p", "360p", "720p"], {"default": "highest"}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("file_path",)
    FUNCTION = "download"
    CATEGORY = "YouTube"

    @staticmethod
    def _get_output_dir() -> str:
        base_output = None
        if folder_paths and hasattr(folder_paths, "get_output_directory"):
            base_output = folder_paths.get_output_directory()
        elif folder_paths and hasattr(folder_paths, "get_temp_directory"):
            base_output = folder_paths.get_temp_directory()
        else:
            base_output = os.path.abspath(os.path.join(os.getcwd(), "output"))
        yt_dir = os.path.join(base_output, "yt_dl")
        os.makedirs(yt_dir, exist_ok=True)
        return yt_dir

    def download(self, url: str, mode: str = "video", quality: str = "highest") -> Tuple[str]:
        import yt_dlp

        output_dir = self._get_output_dir()

        def detect_site(u: str) -> str:
            u = (u or "").lower()
            if "twitter.com" in u or "x.com" in u:
                return "twitter"
            if "youtube.com" in u or "youtu.be" in u:
                return "youtube"
            return "generic"

        site = detect_site(url)

        # Choose format string based on site and mode
        format_string = None
        if site == "youtube":
            if mode == "video":
                format_string = "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"
                if quality != "highest":
                    format_string += f"[height<={quality}]"
            elif mode == "audio":
                format_string = "bestaudio[ext=m4a]/best[ext=m4a]"
        elif site == "twitter":
            if mode == "subtitles":
                # Subtitles generally not available for Twitter; fall back to best video
                print("[yt_dl] Subtitles not supported for Twitter URLs; downloading best video instead.")
                mode = "video"
            if mode == "video":
                # Prefer a direct MP4 if available
                format_string = "best[ext=mp4]/best"
            elif mode == "audio":
                format_string = "bestaudio[ext=m4a]/bestaudio/best"
        else:
            # Generic best as a fallback
            format_string = "best"

        outtmpl = os.path.join(output_dir, "%(title)s.%(ext)s")
        ydl_opts = {
            "format": format_string or "best",
            "outtmpl": outtmpl,
            "writethumbnail": True,
            "noplaylist": True,
            "merge_output_format": "mp4",
            "quiet": False,
        }

        if mode == "subtitles" and site == "youtube":
            ydl_opts.update({
                "skip_download": True,
                "writesubtitles": True,
                "writeautomaticsub": True,
                "subtitleslangs": ["en"],
                "subtitlesformat": "vtt",
            })

        paths = []

        class CaptureHook(yt_dlp.postprocessor.common.PostProcessor):
            def run(self, information):
                filepath = information.get("filepath") or information.get("requested_downloads", [{}])[0].get("filepath")
                if filepath:
                    paths.append(filepath)
                return [], information

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            try:
                ydl.add_post_processor(CaptureHook(ydl))
            except Exception:
                pass
            ydl.download([url])

        file_path = paths[-1] if paths else ""
        return (file_path or "",)


NODE_CLASS_MAPPINGS = {
    "YouTubeDownloader": YouTubeDownloader,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "YouTubeDownloader": "YouTube Download",
}

